package com.clinbrain;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.serializer.SerializerFeature;
import com.facebook.presto.spi.function.Description;
import com.facebook.presto.spi.function.ScalarFunction;
import com.facebook.presto.spi.function.SqlType;
import com.facebook.presto.spi.type.StandardTypes;
import io.airlift.slice.Slice;
import io.airlift.slice.Slices;

import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;

public final class JsonUnfold {

    private JsonUnfold() {

    }

    @ScalarFunction("json_unfold")
    @Description("unfold json, return varchar")
    @SqlType("varchar")
    public static Slice jsonUnfold(@SqlType(StandardTypes.VARCHAR) Slice jsonSlice) {
        Map<String, String> unfoldContent = new HashMap<>();
        if (null != jsonSlice && 0 != jsonSlice.length()) {
            String jsonStr = jsonSlice.toStringUtf8();
            Object parse = JSON.parse(jsonStr);
            LinkedList<String> keys = new LinkedList<>();
            foreachJson(parse, keys, unfoldContent);
        }

        return Slices.utf8Slice(JSON.toJSONString(unfoldContent, SerializerFeature.WriteMapNullValue, SerializerFeature.MapSortField));
    }

    private static void foreachJson(Object object, LinkedList<String> keys, Map<String, String> unfoldContent) {
        if (object instanceof JSONArray) {
            JSONArray jsonArray = (JSONArray) object;
            if (jsonArray.size() > 0) {
                for (int i = 0; i < jsonArray.size(); i++) {
                    keys.add(String.valueOf(i));
                    foreachJson(jsonArray.get(i), keys, unfoldContent);
                    keys.pollLast();
                }
            }
        } else if (object instanceof JSONObject) {
            JSONObject jsonObject = (JSONObject) object;
            Set<String> keySet = jsonObject.keySet();
            Iterator<String> iterator = keySet.iterator();
            while (iterator.hasNext()) {
                String key = iterator.next();
                keys.add(key);

                Object value = jsonObject.get(key);
                if (value instanceof JSONArray) {
                    JSONArray innerArr = (JSONArray) value;
                    foreachJson(innerArr, keys, unfoldContent);
                } else if (value instanceof JSONObject) {
                    foreachJson(value, keys, unfoldContent);
                } else {
                    unfoldContent.put(join(keys), null == value ? null : value.toString());
                }

                keys.pollLast();
            }
        } else {
            unfoldContent.put(join(keys), null == object ? null : object.toString());
        }
    }

    private static String join(List<String> list) {
        StringBuilder keyBuilder = new StringBuilder();
        for (String value : list) {
            keyBuilder.append(value);
            keyBuilder.append(".");
        }

        keyBuilder.deleteCharAt(keyBuilder.length() - 1);
        return keyBuilder.toString();
    }

    public static void main(String[] args) {
        String jsonStr = "{\n" +
                "\t\"name\": \"王小二\",\n" +
                "\t\"age\": 25.2,\n" +
                "\t\"birthday\": \"1990-01-01\",\n" +
                "\t\"school\": \"蓝翔\",\n" +
                "\t\"major\": [\n" +
                "\t\t\"理发\",\n" +
                "\t\t\"挖掘机\"\n" +
                "\t],\n" +
                "\t\"has_girlfriend\": false,\n" +
                "\t\"car\": null,\n" +
                "\t\"house\": null,\n" +
                "\t\"comment\": \"这是一个注释\"\n" +
                "}";
        Slice jsonUnfold = jsonUnfold(Slices.utf8Slice(jsonStr));
        System.out.println(jsonUnfold.toStringUtf8());
    }
}
